#!/usr/bin/env python3
"""
后台任务管理器
用于处理后台包安装任务
"""

import threading
import json
import time
import logging
from datetime import datetime
from typing import Optional, Dict, Any
from system_scripts.package_manager import PackageManager

logger = logging.getLogger('BackgroundTaskManager')

class BackgroundTaskManager:
    """后台任务管理器"""
    
    def __init__(self, db_session, task_id: int):
        self.db_session = db_session
        self.task_id = task_id
        self.task = None
        self.package_manager = None
        self.is_running = False
        self.thread = None
        self.process = None  # 新增：保存pip子进程对象
        
    def load_task(self):
        """加载任务信息"""
        from models import BackgroundInstallationTask
        self.task = self.db_session.get(BackgroundInstallationTask, self.task_id)
        return self.task is not None
    
    def update_task_status(self, status: str, output: str = None, error_message: str = None):
        """更新任务状态"""
        if self.task:
            self.task.status = status
            if output is not None:
                self.task.output = output
            if error_message is not None:
                self.task.error_message = error_message
            if status in ['success', 'failed', 'cancelled']:
                self.task.end_time = datetime.utcnow()
                self.task.duration = self.task.get_elapsed_time()
            self.task.updated_at = datetime.utcnow()
            self.db_session.commit()
    
    def update_progress(self, message: str, level: str = "info"):
        """更新进度信息"""
        if self.task:
            try:
                # 解析现有的进度消息
                if self.task.progress_messages:
                    progress_list = json.loads(self.task.progress_messages)
                else:
                    progress_list = []
                
                # 添加新消息
                progress_list.append({
                    'timestamp': time.time(),
                    'message': message,
                    'level': level
                })
                
                # 保持最新的100条消息
                if len(progress_list) > 100:
                    progress_list = progress_list[-100:]
                
                # 保存到数据库
                self.task.progress_messages = json.dumps(progress_list, ensure_ascii=False)
                self.task.updated_at = datetime.utcnow()
                self.db_session.commit()
                
                logger.info(f"任务 {self.task_id} 进度更新: {message}")
                
            except Exception as e:
                logger.error(f"更新进度失败: {str(e)}")
    
    def progress_callback(self, message: str, level: str = "info"):
        """进度回调函数"""
        self.update_progress(message, level)
    
    def run_installation(self):
        """运行安装任务"""
        from app import create_app
        app = create_app()
        import threading
        import time
        with app.app_context():
            if not self.load_task():
                logger.error(f"任务 {self.task_id} 不存在")
                return
            try:
                self.is_running = True
                self.update_task_status('running')
                logger.info(f"开始后台安装任务 {self.task_id}: {self.task.package_name}")
                if not self.is_running:
                    logger.info(f"任务 {self.task_id} 已被取消")
                    return
                self.package_manager = PackageManager(
                    mirror_source=self.task.mirror_source,
                    progress_callback=self.progress_callback
                )
                # 记录开始时间
                start_time = time.time()
                max_duration = 600  # 10分钟
                # 包装pip命令执行，保存process对象
                def run_with_process(method, *args, **kwargs):
                    # method: install_package/upgrade_package/uninstall_package
                    # monkey patch: 让PackageManager保存process对象
                    orig_popen = self.package_manager._popen if hasattr(self.package_manager, '_popen') else None
                    def my_popen(*p_args, **p_kwargs):
                        self.process = orig_popen(*p_args, **p_kwargs)
                        return self.process
                    if hasattr(self.package_manager, '_popen'):
                        self.package_manager._popen = my_popen
                    result = method(*args, **kwargs)
                    if hasattr(self.package_manager, '_popen'):
                        self.package_manager._popen = orig_popen
                    return result
                # 兼容原有代码，直接在install/upgrade/uninstall里赋值self.process
                # 执行安装/升级/卸载
                if self.task.operation == 'install':
                    success, output, error = self.package_manager.install_package(
                        self.task.package_name, 
                        self.task.version,
                        background_manager=self  # 传递self用于保存process
                    )
                elif self.task.operation == 'upgrade':
                    success, output, error = self.package_manager.upgrade_package(
                        self.task.package_name,
                        background_manager=self
                    )
                elif self.task.operation == 'uninstall':
                    success, output, error = self.package_manager.uninstall_package(
                        self.task.package_name,
                        background_manager=self
                    )
                else:
                    raise ValueError(f"不支持的操作类型: {self.task.operation}")
                # 监控最大运行时长
                while self.is_running and self.process and self.process.poll() is None:
                    if time.time() - start_time > max_duration:
                        logger.warning(f"任务 {self.task_id} 超过最大运行时长，自动终止")
                        self.stop()
                        break
                    time.sleep(1)
                # 检查任务是否被取消
                if not self.is_running:
                    logger.info(f"任务 {self.task_id} 在执行过程中被取消")
                    return
                if success:
                    self.update_task_status('success', output)
                    logger.info(f"后台任务 {self.task_id} 执行成功")
                else:
                    self.update_task_status('failed', output, error)
                    logger.error(f"后台任务 {self.task_id} 执行失败: {error}")
                try:
                    from models import PackageInstallationLog
                    log = PackageInstallationLog(
                        package_name=self.task.package_name,
                        version=self.task.version,
                        operation=(self.task.operation or 'unknown'),  # 保证operation有值
                        status='success' if success else 'failed',
                        output=output,
                        error_message=error,
                        user_id=self.task.user_id,
                        duration=self.task.get_elapsed_time()
                    )
                    self.db_session.add(log)
                    self.db_session.commit()
                    logger.info(f"后台任务 {self.task_id} 日志已保存到包安装日志表")
                except Exception as e:
                    logger.error(f"保存包安装日志失败: {str(e)}")
            except Exception as e:
                error_msg = f"后台任务执行异常: {str(e)}"
                self.update_task_status('failed', error_message=error_msg)
                logger.error(f"后台任务 {self.task_id} 异常: {str(e)}")
            finally:
                self.is_running = False
                self.process = None
    
    def start(self):
        """启动后台任务"""
        if self.thread and self.thread.is_alive():
            logger.warning(f"任务 {self.task_id} 已在运行")
            return False
        
        self.thread = threading.Thread(target=self.run_installation, daemon=True)
        self.thread.start()
        return True
    
    def stop(self):
        """停止后台任务"""
        if self.is_running:
            self.is_running = False
            if self.process and self.process.poll() is None:
                try:
                    self.process.terminate()
                    logger.info(f"已终止pip进程，任务ID: {self.task_id}")
                except Exception as e:
                    logger.error(f"终止pip进程失败: {str(e)}")
            self.update_task_status('cancelled')
            logger.info(f"任务 {self.task_id} 已取消")
            return True
        return False
    
    def get_status(self) -> Dict[str, Any]:
        """获取任务状态"""
        if not self.load_task():
            return {'error': '任务不存在'}
        
        return {
            'id': self.task.id,
            'package_name': self.task.package_name,
            'version': self.task.version,
            'operation': self.task.operation,
            'status': self.task.status,
            'is_running': self.is_running,
            'elapsed_time': self.task.get_elapsed_time(),
            'progress_messages': json.loads(self.task.progress_messages) if self.task.progress_messages else [],
            'output': self.task.output,
            'error_message': self.task.error_message,
            'start_time': self.task.start_time.isoformat() if self.task.start_time else None,
            'end_time': self.task.end_time.isoformat() if self.task.end_time else None
        }

# 全局任务管理器字典
background_tasks = {}

def create_background_task(db_session, user_id: int, package_name: str, operation: str, 
                          version: str = None, mirror_source: str = None) -> int:
    """创建后台任务"""
    from models import BackgroundInstallationTask
    
    # 对于安装操作，检查包是否已安装
    if operation == 'install':
        package_manager = PackageManager()
        if package_manager.check_package_installed(package_name):
            # 获取已安装包的版本信息
            package_info = package_manager.get_package_info(package_name)
            installed_version = package_info['version'] if package_info else '未知版本'
            
            # 无论是否指定版本，只要包已安装就禁止安装
            if version:
                raise ValueError(f'包 {package_name} 已安装（当前版本：{installed_version}），如需安装版本 {version} 请使用升级功能')
            else:
                raise ValueError(f'包 {package_name} 已安装（版本：{installed_version}），如需升级请使用升级功能')
    
    # 创建任务记录
    task = BackgroundInstallationTask(
        user_id=user_id,
        package_name=package_name,
        version=version,
        operation=operation,
        mirror_source=mirror_source,
        status='running'
    )
    
    db_session.add(task)
    db_session.commit()
    
    # 创建任务管理器
    task_manager = BackgroundTaskManager(db_session, task.id)
    background_tasks[task.id] = task_manager
    
    # 启动任务
    task_manager.start()
    
    logger.info(f"创建后台任务 {task.id}: {package_name} {operation}")
    return task.id

def get_background_task(task_id: int) -> Optional[BackgroundTaskManager]:
    """获取后台任务"""
    return background_tasks.get(task_id)

def get_user_running_tasks(user_id: int) -> list:
    """获取用户正在运行的任务"""
    running_tasks = []
    for task_id, task_manager in background_tasks.items():
        if task_manager.task and task_manager.task.user_id == user_id and task_manager.task.is_running():
            running_tasks.append(task_manager.get_status())
    return running_tasks

def cleanup_completed_tasks():
    """清理已完成的任务"""
    completed_task_ids = []
    for task_id, task_manager in background_tasks.items():
        if task_manager.task and task_manager.task.is_completed():
            completed_task_ids.append(task_id)
    
    for task_id in completed_task_ids:
        del background_tasks[task_id]
        logger.info(f"清理已完成任务 {task_id}")

def cleanup_orphaned_tasks():
    """清理孤立的任务（数据库中状态为running但内存中不存在）"""
    try:
        from models import BackgroundInstallationTask
        
        # 获取数据库中所有运行中的任务
        db_running_tasks = BackgroundInstallationTask.query.filter_by(status='running').all()
        
        for db_task in db_running_tasks:
            # 检查内存中是否存在对应的任务管理器
            task_manager = background_tasks.get(db_task.id)
            if not task_manager:
                # 孤立任务，标记为失败
                db_task.status = 'failed'
                db_task.error_message = '任务管理器丢失，标记为失败'
                db_task.end_time = datetime.utcnow()
                db_task.duration = db_task.get_elapsed_time()
                logger.warning(f"清理孤立任务 {db_task.id}: {db_task.package_name}")
        
        # 提交更改
        from app import db
        db.session.commit()
        
    except Exception as e:
        logger.error(f"清理孤立任务失败: {str(e)}") 